import time
import threading
import numpy as np
import pandas as pd
import sqlite3
from tqsdk import TqApi, TqAuth, TqKq
from MyTT import EMA, SUM
from rich.console import Console
from rich.theme import Theme

# 定义颜色主题
custom_theme = Theme({
    "info": "cyan",
    "warning": "yellow",
    "error": "red",
    "success": "green",
    "trade": "bold magenta"
})
console = Console(theme=custom_theme)


class ThreadedTradingBot:
    def __init__(self):
        self.api = None
        self.running = True
        self.threads = []
        self.lock = threading.Lock()
        self.console = console

    def get_monitored_symbols(self):
        """从数据库获取监控品种"""
        conn = sqlite3.connect('trade_monitor.db')
        cursor = conn.cursor()
        cursor.execute("SELECT * FROM products WHERE is_monitored=1")
        products = cursor.fetchall()
        conn.close()
        return [row[2] for row in products]

    def run(self):
        """主运行函数"""
        try:
            self.api = TqApi(account=TqKq(), auth=TqAuth("cps168", "alibaba999"))
            symbols = self.get_monitored_symbols()
            self.console.print(f"开始监控品种: {symbols}", style="info")

            # 为每个品种创建交易线程
            for symbol in symbols:
                thread = threading.Thread(
                    target=self.trade_loop,
                    args=(symbol,),
                    daemon=True
                )
                thread.start()
                self.threads.append(thread)

            # 主线程保持运行
            while self.running:
                time.sleep(1)

        except Exception as e:
            self.console.print(f"程序出错: {e}", style="error")
        finally:
            self.running = False
            for thread in self.threads:
                thread.join()
            if self.api:
                self.api.close()

    def trade_loop(self, symbol):
        """单个品种的交易循环"""
        try:
            quote = self.api.get_quote(symbol)
            klines = self.api.get_kline_serial("SHFE.hc2510", 60,data_length=300)
            #print(klines)
            close = klines.iloc[:].close.values
            # low = klines.iloc[:].low.values
            # open_price = klines.iloc[:].open.values
            # high = klines.iloc[:].high.values
            # print(close,low,open,high)

            # 注册更新通知
            update_event = threading.Event()
            self.api.register_update_notify(update_event.set)

            while self.running:
                # 等待更新事件或超时
                update_event.wait(timeout=1)
                update_event.clear()

                if self.api.is_changing(quote):
                    self.check_stoploss(symbol, quote, klines)

        except Exception as e:
            self.console.print(f"处理品种 {symbol} 时出错: {e}", style="error")

    def check_stoploss(self, symbol, quote, klines):
        """止损逻辑"""
        self.console.print(f"{time.strftime('%X')} {symbol} 价格: {quote.last_price}", style="info")
        try:
            # 假设这些属性是方法，调用它们获取数据
            print(klines)
            close = klines.iloc[:].close.values
            low = klines.iloc[:].low.values
            open_price = klines.iloc[:].open.values
            high = klines.iloc[:].high.values
            print(close,low,open,high)

            q = (3 * close + low + open_price + high) / 6

            # 使用 MyTT 计算加权移动平均线
            weights = np.arange(26, 0, -1)
            wma = np.zeros(len(q))
            for i in range(25, len(q)):
                wma[i] = SUM(q[i - 25:i + 1] * weights, 26) / 351

            # 使用 MyTT 计算 EMA7
            ema_7 = EMA(pd.Series(wma), 7)

            # 使用锁保证线程安全
            with self.lock:
                position = self.api.get_position(symbol)

                # 平仓条件判断
                if len(wma) >= 2 and len(ema_7) >= 2:
                    if position.pos_short > 0 and wma[-2] > ema_7[-2]:
                        self.api.insert_order(
                            symbol=symbol,
                            direction="BUY",
                            offset="CLOSE",
                            volume=position.pos_short
                        )
                        self.console.print(
                            f"{time.strftime('%X')} 平空仓: {symbol} 手数: {position.pos_short}",
                            style="trade"
                        )

                    if position.pos_long > 0 and wma[-2] < ema_7[-2]:
                        self.api.insert_order(
                            symbol=symbol,
                            direction="SELL",
                            offset="CLOSE",
                            volume=position.pos_long
                        )
                        self.console.print(
                            f"{time.strftime('%X')} 平多仓: {symbol} 手数: {position.pos_long}",
                            style="trade"
                        )

        except Exception as e:
            self.console.print(f"处理 {symbol} 止损逻辑时出错: {e}", style="error")


if __name__ == "__main__":
    bot = ThreadedTradingBot()
    try:
        bot.run()
    except KeyboardInterrupt:
        bot.running = False
        console.print("\n程序已停止", style="warning")


